import logging
import argparse
import os
import json

def load_video_paths(annotation_file, video_name_mapping, video_dir):
    """
    返回视频路径列表，每个元素是一个元组，元组中包含了视频的路径和视频的id
    例如：[('/root/autodl-tmp/MSVD-QA/video/00jrXRMlZOY_0_10.avi', 47)]
    :param annotation_file: annotation 文件路径
    :param video_name_mapping: video name mapping 文件路径
    :param video_dir: 视频文件夹路径
    :return: video_paths: 视频路径列表
    :author: njupt-mcc(github), qms
    """
    video_paths = []
    video_ids = []
    with open(annotation_file, 'r') as anno_file:  # MSVD-QA/train_qa.json
        instances = json.load(anno_file)
        # {"answer":"animal","id":0,"question":"what is chewing on a nut?","video_id":1}
        [video_ids.append(instance['video_id']) for instance in instances]
    video_ids = set(video_ids)
    with open(video_name_mapping, 'r') as mapping:
        mapping_pairs = mapping.read().split('\n')  # 读取文件中的每一行
        # 00jrXRMlZOY_0_10 vid47
    mapping_dict = {}
    for idx in range(len(mapping_pairs)):
        cur_pair = mapping_pairs[idx].split(' ')  # 读取每一行中的两个元素
        mapping_dict[cur_pair[1]] = cur_pair[0]  # vid47: 00jrXRMlZOY_0_10
    for video_id in video_ids:
        video_paths.append(
            (video_dir + '/{}.avi'.format(mapping_dict['vid' + str(video_id)]), video_id)
            # '/root/autodl-tmp/MSVD-QA/video' + '{}.avi'.format(00jrXRMlZOY_0_10), 47
            # /root/autodl-tmp/MSVD-QA/video/00jrXRMlZOY_0_10.avi, 47
        )
    return video_paths

def load_video_id_to_video_path(video_paths):
    """
    将video_paths转换成字典，键是video_id，值是video_path
    :param video_paths: 视频路径列表，每个元素是一个元组，元组中包含了视频的路径和视频的id
    :return: video_id_to_video_path: video_id到video_path的映射
    :author: qms
    """
    video_id_to_video_path = {}
    for video_path in video_paths:
        video_id_to_video_path[video_path[1]] = video_path[0]
    return video_id_to_video_path

logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)-8s %(message)s')
parser = argparse.ArgumentParser()
parser.add_argument('--annotation_file', type=str, default='data/msvd-qa/train_qa.json')
parser.add_argument('--video_dir', type=str, default='data/msvd-qa/video')
parser.add_argument('--video_name_mapping', type=str, default='data/msvd-qa/youtube_mapping.txt')
opt = parser.parse_args()

video_paths = load_video_paths(opt.annotation_file, opt.video_name_mapping, opt.video_dir)
video_id_to_video_path = load_video_id_to_video_path(video_paths)
for video_id, path in video_id_to_video_path.items():
    if not os.path.exists(path):
        logging.info('video_id: %s, path: %s' % (video_id, path))
